{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Feature Space Illumination with MAPElites\n",
    "\n",
    "`MAPElites` (stylized as MAP-Elites in the original paper [1]) is an algorithm where the feature space (or sometimes called the behavior space) is organized into cells, and each cell has its own local competition: if a new solution's feature(s) is/are suitable for a cell and this new solution has a better fitness than the cell's current solution, the new solution replaces the old one.\n",
    "Organizing the population this way forces diversity, and therefore allows one to illuminate the search space according to the criteria specified by the user.\n",
    "\n",
    "In this notebook, we use `MAPElites` to illuminate the search space of the bi-objective test problem of Kursawe [2]. Given these two functions:\n",
    "\n",
    "$$\n",
    "\\begin{array}{r c l}\n",
    "f_1(x) &=& \\sum_{i = \\{1, 2\\}} \\Big[-10 \\cdot \\text{exp} \\big(-0.2 \\cdot \\sqrt{x^2_i + x^2_{i+1}} \\big)\\Big] \\\\\n",
    "f_2(x) &=& \\sum_{i = \\{1, 2, 3\\}} \\Big[|x_i|^{0.8} + 5 \\text{sin}\\big(x^3_i\\big)\\Big]\n",
    "\\end{array}\n",
    "$$\n",
    "\n",
    "Kursawe's test problem, in its original form, is defined as:\n",
    "\n",
    "$$\n",
    "\\begin{array}{r l}\n",
    "\\text{minimize} & f_1(x) \\\\\n",
    "\\text{minimize} & f_2(x) \\\\\n",
    "\\text{subject to} & -5 \\leq x_i \\leq 5 \\quad \\forall i \\in \\{1, 2, 3\\}\n",
    "\\end{array}\n",
    "$$\n",
    "\n",
    "In this notebook, however, we modify this problem such that the goal is to minimize $f_1(x) + f_2(x)$, and the two features are $f_1(x)$ and $f_2(x)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from evotorch import Problem\n",
    "from evotorch.algorithms import MAPElites\n",
    "from evotorch.operators import GaussianMutation\n",
    "from evotorch.logging import StdOutLogger, PandasLogger\n",
    "import matplotlib.pyplot as plt\n",
    "import torch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "Below is the definition of the bi-objective Kursawe function. The input of this function is a tensor `x` of shape `(n, 3)`, `n` being the number of solutions. The returned tensor is of shape `(n, 2)`, in which the column with index 0 stores $f_1(\\cdot)$ and column with index 1 stores $f_2(\\cdot)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def kursawe(x: torch.Tensor) -> torch.Tensor:\n",
    "    f1 = torch.sum(\n",
    "        -10 * torch.exp(\n",
    "            -0.2 * torch.sqrt(x[:, 0:2] ** 2.0 + x[:, 1:3] ** 2.0)\n",
    "        ),\n",
    "        dim=-1,\n",
    "    )\n",
    "    f2 = torch.sum(\n",
    "        (torch.abs(x) ** 0.8) + (5 * torch.sin(x ** 3)),\n",
    "        dim=-1,\n",
    "    )\n",
    "    fitnesses = torch.stack([f1, f2], dim=-1)\n",
    "    return fitnesses"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "Below is a wrapper for the kursawe function. The return value of the wrapped/modified Kursawe function is a tensor of shape `(n, 3)` where the column with index 0 stores $f_1(\\cdot)+f_2(\\cdot)$, column with index 1 stores $f_1(\\cdot)$, and column with index 2 stores $f_2(\\cdot)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def kursawe_with_features(x: torch.Tensor) -> torch.Tensor:\n",
    "    fitnesses = kursawe(x)\n",
    "    total_fitness = torch.sum(fitnesses, dim=-1).reshape(-1, 1)\n",
    "    return torch.hstack([total_fitness, fitnesses])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "We now define our optimization problem. Notice that:\n",
    "\n",
    "- We have only one objective with objective sense \"min\"\n",
    "- We have extra evaluation data whose length is 2 (`eval_data_length=2`)\n",
    "\n",
    "With this configuration, we declare that we expect evaluation result tensors to have 3 columns: leftmost column for the objective, and the remaining columns for features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "problem = Problem(\n",
    "    \"min\",\n",
    "    kursawe_with_features,\n",
    "    solution_length=3,\n",
    "    eval_data_length=2,\n",
    "    bounds=(-5.0, 5.0),\n",
    "    vectorized=True,\n",
    ")\n",
    "\n",
    "problem"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "Below, we create a hypergrid for our feature space. In our hypergrid, the global lower bound for each feature is -20, and the global upper bound for each feature is 20. For each feature, we declare that we want 50 bins."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_grid = MAPElites.make_feature_grid(\n",
    "    lower_bounds=[-20, -14],\n",
    "    upper_bounds=[-10, 4],\n",
    "    num_bins=50,\n",
    "    dtype=\"float32\",\n",
    ")\n",
    "\n",
    "feature_grid.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "Now that we have our hypergrid, we can instantiate `MAPElites`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "searcher = MAPElites(\n",
    "    problem,\n",
    "    operators=[GaussianMutation(problem, stdev=0.1)],\n",
    "    re_evaluate=False,\n",
    "    feature_grid=feature_grid,\n",
    ")\n",
    "\n",
    "searcher"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "Now we run our evolutionary search."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "StdOutLogger(searcher, interval=10)\n",
    "pandas_logger = PandasLogger(searcher)\n",
    "\n",
    "searcher.run(100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "Below we can see how many of the cells within our hypergrid are filled:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.count_nonzero(searcher.filled)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "Now we plot our solutions.\n",
    "In the plot, dark colors represent better solutions. The darkest solutions, therefore, should give us an outline of the pareto-front."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_f1 = []\n",
    "all_f2 = []\n",
    "all_sumf = []\n",
    "\n",
    "for i, solution in enumerate(searcher.population):\n",
    "    if searcher.filled[i]:\n",
    "        f1 = float(solution.evals[1])\n",
    "        f2 = float(solution.evals[2])\n",
    "        all_f1.append(f1)\n",
    "        all_f2.append(f2)\n",
    "        all_sumf.append(f1 + f2)\n",
    "\n",
    "min_sumf = min(all_sumf)\n",
    "max_sumf = max(all_sumf)\n",
    "diff_sumf = max_sumf - min_sumf\n",
    "\n",
    "c = []\n",
    "for i in range(len(all_f1)):\n",
    "    norm_f = (all_sumf[i] - min_sumf) / (max_sumf - min_sumf)\n",
    "    color = hex(int(norm_f * 255))[2:]\n",
    "    if len(color) == 1:\n",
    "        color = \"0\" + color\n",
    "    color = f\"#{color}{color}{color}\"\n",
    "    c.append(color)\n",
    "\n",
    "plt.scatter(all_f1, all_f2, c=c)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "# References\n",
    "\n",
    "[1] Mouret, J. B., & Clune, J. (2015). Illuminating search spaces by mapping elites. arXiv preprint arXiv:1504.04909.\n",
    "\n",
    "[2] Kursawe, F. (1991). A variant of evolution strategies for vector optimization. In: Schwefel, HP., Männer, R. (eds) Parallel Problem Solving from Nature. PPSN 1990. Lecture Notes in Computer Science, vol 496. Springer, Berlin, Heidelberg. https://doi.org/10.1007/BFb0029752\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
